#include "utils.h"
#include <assert.h>
#include <cstdint>
#include <functional>
#include <limits>
#include <memory>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <string>
#include <unistd.h>
#include <vector>

// Exception class. Caught in main().
class Error : public std::exception {
  std::string msg_;

public:
  Error() = delete; // No default constructor.
  explicit Error(const char *msg) : msg_(msg) {}
  explicit Error(std::string &&msg) : msg_(std::move(msg)) {}

  const char *what() const noexcept override { return msg_.c_str(); }
};

void write_hexdigit(WriteBuf &buf, const uint8_t x) {
  if (x < 10) {
    buf.write_uint8('0' + x);
  } else if (x < 16) {
    buf.write_uint8('A' + x - 10);
  } else {
    throw Error("Bad hex digit");
  }
}

void write_octal_uint8(WriteBuf &buf, const uint8_t x) {
  write_hexdigit(buf, x >> 6);
  write_hexdigit(buf, (x >> 3) & 0x7);
  write_hexdigit(buf, x & 0x7);
}

void write_hex_uint8(WriteBuf &buf, const uint8_t x) {
  write_hexdigit(buf, x >> 6);
  write_hexdigit(buf, (x >> 3) & 0x7);
  write_hexdigit(buf, x & 0x7);
}

void write_stringobj(WriteBuf &buf, const std::string &str) {
  buf.write_string("(");
  for (auto c : str) {
    if (isprint(c)) {
      buf.write_uint8(c);
    } else {
      buf.write_uint8('\\');
      write_octal_uint8(buf, static_cast<uint8_t>(c));
    }
  }
  buf.write_string(")");
}

void write_nameobj(WriteBuf &buf, const std::string &str) {
  buf.write_string("/");
  for (auto c : str) {
    if (isprint(c)) {
      buf.write_uint8(c);
    } else {
      buf.write_uint8('#');
      write_octal_uint8(buf, static_cast<uint8_t>(c));
    }
  }
  buf.write_string(" ");
}

void write_intobj(WriteBuf &buf, int i) {
  char str[32];
  snprintf(str, sizeof(str), "%d ", i);
  buf.write_string(str);
}

void write_numobj(WriteBuf &buf, double d) {
  char str[64];
  snprintf(str, sizeof(str), "%f ", d);
  buf.write_string(str);
}

void write_command(WriteBuf &buf, const std::string &cmd) {
  buf.write_string(cmd.c_str());
  buf.write_string("\n");
}

class PDF {
public:
  PDF() {}
  virtual ~PDF() {}

  virtual void write(WriteBuf &buf) const = 0;
};

typedef std::unique_ptr<PDF> PDFptr;
typedef std::vector<PDFptr> PDFvec;

// Utility for reading the current file offset.
class PDF_ReadPreOffset : public PDF {
  const std::function<void(size_t)> f_;
  const PDFptr child_;

public:
  PDF_ReadPreOffset(std::function<void(size_t)> &&f, PDFptr &&child)
      : f_(std::move(f)), child_(std::move(child)) {}

  static std::unique_ptr<PDF_ReadPreOffset> mk(std::function<void(size_t)> &&f,
                                               PDFptr &&child) {
    return std::make_unique<PDF_ReadPreOffset>(std::move(f), std::move(child));
  }

  void write(WriteBuf &buf) const override {
    f_(buf.offset());
    child_->write(buf);
  }
};

// Utility for reading the current file offset.
class PDF_ReadPostOffset : public PDF {
  const std::function<void(size_t)> f_;
  const PDFptr child_;

public:
  PDF_ReadPostOffset(std::function<void(size_t)> &&f, PDFptr &&child)
      : f_(std::move(f)), child_(std::move(child)) {}

  static std::unique_ptr<PDF_ReadPostOffset> mk(std::function<void(size_t)> &&f,
                                                PDFptr &&child) {
    return std::make_unique<PDF_ReadPostOffset>(std::move(f), std::move(child));
  }

  void write(WriteBuf &buf) const override {
    child_->write(buf);
    f_(buf.offset());
  }
};

class PDF_Int : public PDF {
  const int i_;

public:
  explicit PDF_Int(int i) : i_(i) {}

  static std::unique_ptr<PDF_Int> mk(int i) {
    return std::make_unique<PDF_Int>(i);
  }

  void write(WriteBuf &buf) const override { write_intobj(buf, i_); }
};

// Like PDF_Int, except with padded output so that the number of
// characters is always the same. This is useful for integer values
// that aren't known until the second pass.
class PDF_IntF : public PDF {
  const int i_;
  const int w_;

public:
  explicit PDF_IntF(int i, int w) : i_(i), w_(w) {}

  static std::unique_ptr<PDF_IntF> mk(int i, int w = 10) {
    return std::make_unique<PDF_IntF>(i, w);
  }

  void write(WriteBuf &buf) const override {
    char str[32];
    assert(0 <= w_ && w_ < static_cast<int>(sizeof(str)));
    snprintf(str, sizeof(str), "%*d", w_, i_);
    buf.write_bytes((const uint8_t *)str, w_);
    buf.write_string("\n");
  }
};

class PDF_Num : public PDF {
  const double d_;

public:
  explicit PDF_Num(double d) : d_(d) {}

  static std::unique_ptr<PDF_Num> mk(double d) {
    return std::make_unique<PDF_Num>(d);
  }

  void write(WriteBuf &buf) const override { write_numobj(buf, d_); }
};

class PDF_Ref : public PDF {
  const int num_;
  const int gen_;

public:
  PDF_Ref(int num, int gen) : num_(num), gen_(gen) {}

  static std::unique_ptr<PDF_Ref> mk(int num, int gen) {
    return std::make_unique<PDF_Ref>(num, gen);
  }

  void write(WriteBuf &buf) const override {
    write_intobj(buf, num_);
    write_intobj(buf, gen_);
    buf.write_string("R ");
  }
};

class PDF_Cmd : public PDF {
  const std::string cmd_;

public:
  explicit PDF_Cmd(std::string &&cmd) : cmd_(std::move(cmd)) {}

  static std::unique_ptr<PDF_Cmd> mk(std::string &&cmd) {
    return std::make_unique<PDF_Cmd>(std::move(cmd));
  }

  void write(WriteBuf &buf) const override { write_command(buf, cmd_); }
};

class PDF_Name : public PDF {
  const std::string name_;

public:
  explicit PDF_Name(std::string &&name) : name_(std::move(name)) {}

  static std::unique_ptr<PDF_Name> mk(std::string &&name) {
    return std::make_unique<PDF_Name>(std::move(name));
  }

  void write(WriteBuf &buf) const override { write_nameobj(buf, name_); }
};

class PDF_String : public PDF {
  const std::string str_;

public:
  explicit PDF_String(std::string &&str) : str_(std::move(str)) {}

  static std::unique_ptr<PDF_String> mk(std::string &&str) {
    return std::make_unique<PDF_String>(std::move(str));
  }

  void write(WriteBuf &buf) const override { write_stringobj(buf, str_); }
};

class PDF_Comment : public PDF {
  const std::string comment_;

public:
  explicit PDF_Comment(std::string &&comment) : comment_(std::move(comment)) {}

  static std::unique_ptr<PDF_Comment> mk(std::string &&comment) {
    return std::make_unique<PDF_Comment>(std::move(comment));
  }

  void write(WriteBuf &buf) const override {
    buf.write_string("%");
    buf.write_string(comment_.c_str());
    buf.write_string("\n");
  }
};

class PDF_Seq : public PDF {
  const PDFvec seq_;

public:
  explicit PDF_Seq(std::vector<std::unique_ptr<PDF>> &&seq)
      : seq_(std::move(seq)) {}

  static std::unique_ptr<PDF_Seq> mk(PDFvec &&seq) {
    return std::make_unique<PDF_Seq>(std::move(seq));
  }

  void write(WriteBuf &buf) const override {
    for (auto &x : seq_) {
      x->write(buf);
    }
  }
};

class PDF_Array : public PDF {
  const PDFvec array_;

public:
  explicit PDF_Array(std::vector<std::unique_ptr<PDF>> &&array)
      : array_(std::move(array)) {}

  static std::unique_ptr<PDF_Array> mk(PDFvec &&array) {
    return std::make_unique<PDF_Array>(std::move(array));
  }

  void write(WriteBuf &buf) const override {
    buf.write_string("[ ");
    for (auto &x : array_) {
      x->write(buf);
    }
    buf.write_string("] ");
  }
};

// key-value pair for a dict.
struct PDF_KV {
  std::string key_;
  PDFptr value_;

  PDF_KV(std::string &&key, PDFptr &&value)
      : key_(std::move(key)), value_(std::move(value)) {}
};

class PDF_Dict : public PDF {
  const std::vector<PDF_KV> dict_;

public:
  explicit PDF_Dict(std::vector<PDF_KV> &&dict) : dict_(std::move(dict)) {}

  static std::unique_ptr<PDF_Dict> mk(std::vector<PDF_KV> &&dict) {
    return std::make_unique<PDF_Dict>(std::move(dict));
  }

  void write(WriteBuf &buf) const override {
    buf.write_string("<< ");
    for (auto &kv : dict_) {
      write_nameobj(buf, kv.key_);
      kv.value_->write(buf);
    }
    buf.write_string(">> ");
  }
};

class PDF_Stream : public PDF {
  const std::vector<uint8_t> stream_;

public:
  explicit PDF_Stream(const std::string &str)
      : stream_(str.begin(), str.end()) {}

  explicit PDF_Stream(std::vector<uint8_t> &&stream)
      : stream_(std::move(stream)) {}

  static std::unique_ptr<PDF_Stream> mk(const std::string &str) {
    return std::make_unique<PDF_Stream>(str);
  }

  static std::unique_ptr<PDF_Stream> mk(std::vector<uint8_t> &&stream) {
    return std::make_unique<PDF_Stream>(std::move(stream));
  }

  void write(WriteBuf &buf) const override {
    buf.write_string("stream\n");
    buf.write_bytes(stream_.data(), stream_.size());
    buf.write_string(" endstream\n");
  }
};

struct OffsetsTable {
  size_t filesize_ = 0;
  size_t startbody_ = 0;
  size_t startxref_ = 0;

  size_t ref001_ = 0;
  size_t ref002_ = 0;
  size_t ref003_ = 0;
  size_t ref004_ = 0;
  size_t ref005_ = 0;
  size_t ref006_ = 0;

  size_t stream000_start_ = 0;
  size_t stream000_end_ = 0;
};

static const size_t XRefEntrySize = sizeof(uint32_t) + 2 * sizeof(uint64_t);

static void writeXRefEntry(uint8_t *entry, uint32_t type, uint64_t offset,
                           uint64_t gen) {
  *(uint32_t *)entry = htobe32(type);
  entry += sizeof(type);
  *(uint64_t *)entry = htobe64(offset);
  entry += sizeof(offset);
  *(uint64_t *)entry = htobe64(gen);
}

static PDFptr mkAnnotOverflow(const OffsetsTable &offsets) {
  const int first0 = 0;
  const int len0 = 7;
  const int first1 = static_cast<int>(offsets.ref005_);
  const int len1 = 1;
  const int numEntries = len0 + len1;
  const int streamsize = numEntries * XRefEntrySize;
  std::vector<uint8_t> stream(streamsize);
  uint8_t *streamdata = stream.data();

  writeXRefEntry(streamdata, 1, 0, 0);
  streamdata += XRefEntrySize;
  writeXRefEntry(streamdata, 1, offsets.ref001_, 0);
  streamdata += XRefEntrySize;
  writeXRefEntry(streamdata, 1, offsets.ref002_, 0);
  streamdata += XRefEntrySize;
  writeXRefEntry(streamdata, 1, offsets.ref003_, 0);
  streamdata += XRefEntrySize;
  writeXRefEntry(streamdata, 1, offsets.ref004_, 0);
  streamdata += XRefEntrySize;
  writeXRefEntry(streamdata, 2, offsets.ref005_, 0);
  streamdata += XRefEntrySize;
  writeXRefEntry(streamdata, 1, offsets.ref006_, 0);
  streamdata += XRefEntrySize;
  writeXRefEntry(streamdata, 1, offsets.ref005_, 0);

  return PDF_Seq::mk(_vec<PDFptr>(
      PDF_Int::mk(1337), PDF_Int::mk(133713), PDF_Cmd::mk("obj"),
      PDF_Dict::mk(_vec<PDF_KV>(
          PDF_KV(std::string("Size"), PDF_Int::mk(1337)),
          PDF_KV(std::string("Length"), PDF_Int::mk(streamsize)),
          PDF_KV(std::string("Prev"),
                 PDF_Int::mk(-1)), // link to next XRef table
          PDF_KV(std::string("Root"), PDF_Ref::mk(1, 0)),
          PDF_KV(std::string("Index"),
                 PDF_Array::mk(_vec<PDFptr>(
                     PDF_IntF::mk(first0, 4), PDF_IntF::mk(len0, 3),
                     PDF_IntF::mk(first1, 4), PDF_IntF::mk(len1, 3)))),
          PDF_KV(std::string("W"),
                 PDF_Array::mk(_vec<PDFptr>(PDF_Int::mk(4), PDF_Int::mk(8),
                                            PDF_Int::mk(8)))))),
      PDF_Stream::mk(std::move(stream))));
}

static PDFptr mkContents(OffsetsTable &offsets) {
  const int streamsize =
      static_cast<int>(offsets.stream000_end_ - offsets.stream000_start_);
  return PDF_Seq::mk(_vec<PDFptr>(
      PDF_Dict::mk(_vec<PDF_KV>(
          PDF_KV(std::string("Length"), PDF_IntF::mk(streamsize)))),
      PDF_ReadPostOffset::mk(
          [&offsets](size_t offset) { offsets.stream000_start_ = offset; },
          PDF_Cmd::mk("stream")),
      PDF_Name::mk("kevstatearg"), PDF_Cmd::mk("gs"),
      PDF_ReadPreOffset::mk(
          [&offsets](size_t offset) { offsets.stream000_end_ = offset; },
          PDF_Cmd::mk("endstream"))));
}

static const int mkAnnotArray_first = 20;

static std::vector<uint8_t> mkAnnotArray_ObjectStream(size_t nElements,
                                                      size_t nLayers) {
  std::vector<uint8_t> txt(mkAnnotArray_first);

  PDFptr prologue = PDF_Seq::mk(
      _vec<PDFptr>(PDF_IntF::mk(5), PDF_Int::mk(mkAnnotArray_first)));
  WriteBuf buf(txt.data(), mkAnnotArray_first);
  prologue->write(buf);
  while (buf.offset() < mkAnnotArray_first) {
    buf.write_string(" ");
  }

  char reftxt[6] = {'6', ' ', '0', ' ', 'R', ' '};

  const size_t offset = txt.size();
  txt.resize(offset + nElements * sizeof(reftxt) + 2);
  uint8_t *p = txt.data() + offset;
  *p++ = '[';
  for (size_t i = 0; i < nElements; i++) {
    memcpy(p, reftxt, sizeof(reftxt));
    p += sizeof(reftxt);
  }
  *p++ = ']';

  for (size_t i = 0; i < nLayers; i++) {
    std::vector<uint8_t> tmp;
    compress(tmp, txt);
    txt = std::move(tmp);
  }

  return txt;
}

static PDFptr mkAnnots() {
  static const std::vector<uint8_t> annots_txt(
      mkAnnotArray_ObjectStream(0x1000000, 2));
  std::vector<uint8_t> txt = annots_txt;
  const int streamsize = static_cast<int>(txt.size());
  return PDF_Seq::mk(_vec<PDFptr>(
      PDF_Dict::mk(_vec<PDF_KV>(
          PDF_KV(std::string("N"), PDF_IntF::mk(1)),
          PDF_KV(std::string("First"), PDF_Int::mk(mkAnnotArray_first)),
          PDF_KV(std::string("Length"), PDF_Int::mk(streamsize)),
          PDF_KV(std::string("Filter"),
                 PDF_Array::mk(_vec<PDFptr>(PDF_Name::mk("FlateDecode"),
                                            PDF_Name::mk("FlateDecode")))))),
      PDF_Stream::mk(std::move(txt))));
}

static PDFptr mkBody(OffsetsTable &offsets) {
  const int numKids = 256;
  std::vector<PDFptr> kids;
  for (size_t i = 0; i < numKids; i++) {
    kids.push_back(PDF_Ref::mk(3, 0));
  }
  return PDF_Seq::mk(_vec<PDFptr>(
      // root
      PDF_ReadPreOffset::mk(
          [&offsets](size_t offset) { offsets.ref001_ = offset; },
          PDF_Int::mk(1)),
      PDF_Int::mk(0), PDF_Cmd::mk("obj"),
      PDF_Dict::mk(_vec<PDF_KV>(
          PDF_KV(std::string("Pages"), PDF_Ref::mk(2, 0)),
          PDF_KV(std::string("AcroForm"),
                 PDF_Dict::mk(_vec<PDF_KV>(
                     PDF_KV(std::string("Fields"),
                            PDF_Array::mk(_vec<PDFptr>(PDF_Ref::mk(6, 0))))))

                     ),
          PDF_KV(std::string("Size"), PDF_Int::mk(2)),
          PDF_KV(std::string("Length"), PDF_Int::mk(1337)))),
      // pages
      PDF_ReadPreOffset::mk(
          [&offsets](size_t offset) { offsets.ref002_ = offset; },
          PDF_Int::mk(2)),
      PDF_Int::mk(0), PDF_Cmd::mk("obj"),
      PDF_Dict::mk(_vec<PDF_KV>(
          PDF_KV(std::string("Pages"),
                 PDF_Dict::mk(_vec<PDF_KV>(
                     PDF_KV(std::string("Count"), PDF_Int::mk(1))))),
          PDF_KV(std::string("Size"), PDF_Int::mk(2)),
          PDF_KV(std::string("Count"), PDF_Int::mk(numKids)),
          PDF_KV(std::string("Kids"), PDF_Array::mk(std::move(kids))),
          PDF_KV(std::string("Length"), PDF_Int::mk(1337)))),
      // kid
      PDF_ReadPreOffset::mk(
          [&offsets](size_t offset) { offsets.ref003_ = offset; },
          PDF_Int::mk(3)),
      PDF_Int::mk(0), PDF_Cmd::mk("obj"),
      PDF_Dict::mk(
          _vec<PDF_KV>(PDF_KV(std::string("Type"), PDF_Name::mk("Page")),
                       PDF_KV(std::string("Contents"), PDF_Ref::mk(4, 0)),
                       PDF_KV(std::string("Annots"), PDF_Ref::mk(5, 0)),
                       PDF_KV(std::string("Size"), PDF_Int::mk(2)),
                       PDF_KV(std::string("Count"), PDF_Int::mk(1)))),
      // contents
      PDF_ReadPreOffset::mk(
          [&offsets](size_t offset) { offsets.ref004_ = offset; },
          PDF_Int::mk(4)),
      PDF_Int::mk(0), PDF_Cmd::mk("obj"), mkContents(offsets),
      // annots
      PDF_ReadPreOffset::mk(
          [&offsets](size_t offset) { offsets.ref005_ = offset; },
          PDF_IntF::mk(static_cast<int>(offsets.ref005_))),
      PDF_Int::mk(0), PDF_Cmd::mk("obj"), mkAnnots(),
      // widget
      PDF_ReadPreOffset::mk(
          [&offsets](size_t offset) { offsets.ref006_ = offset; },
          PDF_IntF::mk(6)),
      PDF_Int::mk(0), PDF_Cmd::mk("obj"),
      PDF_Dict::mk(_vec<PDF_KV>(
          PDF_KV(std::string("Subtype"), PDF_Name::mk("Widget")),
          PDF_KV(std::string("DV"), PDF_String::mk("kevwozere001")),
          PDF_KV(std::string("V"), PDF_String::mk("kevwozere002")),
          PDF_KV(std::string("Ff"), PDF_Int::mk(0xDEADBEEF)),
          PDF_KV(std::string("MaxLen"), PDF_Int::mk(0xDEADBEEF)),
          PDF_KV(std::string("F"), PDF_Int::mk(2)), // Annot::flagHidden
          PDF_KV(std::string("Rect"),
                 PDF_Array::mk(_vec<PDFptr>(PDF_Int::mk(2), PDF_Int::mk(3),
                                            PDF_Int::mk(4), PDF_Int::mk(5)))),
          PDF_KV(std::string("FT"), PDF_Name::mk("Tx"))))));
}

static PDFptr mkXRef(const OffsetsTable &offsets) {
  return mkAnnotOverflow(offsets);
}

static PDFptr mkPDF(OffsetsTable &offsets) {
  return PDF_Seq::mk(_vec<PDFptr>(
      PDF_Comment::mk("PDF-1.7"), PDF_Int::mk(4), PDF_Int::mk(0),
      PDF_Cmd::mk("obj"),
      PDF_Dict::mk(_vec<PDF_KV>(
          PDF_KV(std::string("Linearized"), PDF_Int::mk(-1)),
          PDF_KV(std::string("L"),
                 PDF_IntF::mk(static_cast<int>(offsets.filesize_))),
          PDF_KV(std::string("T"), PDF_Int::mk(1000000)))),
      PDF_Cmd::mk("endobj"),
      PDF_ReadPreOffset::mk(
          [&offsets](size_t offset) { offsets.startbody_ = offset; },
          mkBody(offsets)),
      PDF_ReadPreOffset::mk(
          [&offsets](size_t offset) { offsets.startxref_ = offset; },
          mkXRef(offsets)),
      PDF_Cmd::mk("startxref"),
      PDF_IntF::mk(static_cast<int>(offsets.startxref_)),
      PDF_ReadPostOffset::mk(
          [&offsets](size_t offset) { offsets.filesize_ = offset; },
          PDF_Comment::mk("%EOF"))));
}

int main() {
  try {
    std::vector<uint8_t> rawbuf(0x10000);

    OffsetsTable offsets;

    WriteBuf buf(rawbuf.data(), rawbuf.size());

    offsets.ref005_ = 100;
    // Two passes. The first pass calculates the values of filesize and
    // startxref.
    PDFptr pdf = mkPDF(offsets);
    pdf->write(buf);

    const size_t oldfilesize = buf.offset();
    buf.reset();

    pdf = mkPDF(offsets);
    pdf->write(buf);
    if (oldfilesize != buf.offset()) {
      throw Error("filesize changed on second pass");
    }

    buf.write_to_fd(STDOUT_FILENO);

    return EXIT_SUCCESS;
  } catch (Error &e) {
    fprintf(stderr, "%s\n", e.what());
    return EXIT_FAILURE;
  }
}
